import re
from datetime import datetime

def extract_year(text):
    """从文本中提取年份"""
    # 搜索第一个匹配的位置
    match = re.search(r'(\d{4})款', text)
    if match:
        # 返回第一个转换为int后捕获的内容
        return int(match.group(1))
    return None

def extract_mileage(text):
    """从文本中提取里程数"""
    # 尝试匹配"xx万公里"格式
    match = re.search(r'(\d+\.?\d*)[万]?公里', text)
    if match:
        mileage = float(match.group(1))
        # 包含万字直接返回，否则除以10000转换为公里
        return mileage if '万' in text else mileage/10000
    return None

def process_data(df):
    """处理原始数据"""
    # 创建结果DataFrame
    result_df = df.copy()
    
    # 使用apply对每行数据进行处理
    # 将文本按行分割并取第一个元素作为车型
    result_df['车型'] = result_df['文本内容'].apply(lambda x: x.split(' ')[0] if isinstance(x, str) else '未知')
    
    # 提取年份并计算车龄
    current_year = datetime.now().year
    result_df['生产年份'] = result_df['文本内容'].apply(extract_year)
    result_df['车龄'] = result_df['生产年份'].apply(lambda x: current_year - x if x else None)
    
    # 提取里程数
    result_df['里程数'] = result_df['文本内容'].apply(extract_mileage)
    
    # 删除缺失值
    result_df = result_df.dropna(subset=['车龄', '里程数'])
    
    # 确保数值类型正确
    result_df['车龄'] = result_df['车龄'].astype(float)
    result_df['里程数'] = result_df['里程数'].astype(float)
    
    print(f"数据处理完成，共 {len(result_df)} 条有效记录")
    print("列名:", result_df.columns.tolist())
    
    return result_df